from sklearn.model_selection import  train_test_split
import  torch
import  pandas as pd
from torch import nn
from tensorboardX import SummaryWriter

import  torch.nn.functional as F
data=pd.read_csv('HR.csv')

print(data.info)
print(data.part.unique())
print(data.salary.unique())
print(data.groupby(['salary','part']).size())

'''独热编码'''
print(pd.get_dummies(data.salary))
data=data.join(pd.get_dummies((data.salary)))
del data['salary']
data=data.join(pd.get_dummies((data.part)))
del data['part']
print("预测出来的准确率必须高于{}才有意义".format(11428/len(data)))

'''数据预处理'''
Y_data=data.left.values.reshape(-1,1)
Y=torch.from_numpy(Y_data).type(torch.float32)
X_data=data[[c for  c in data.columns if c!='left']].values
X=torch.from_numpy(X_data).type(torch.float32)

'''自定义创建模型'''
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.liner_1=nn.Linear(20,64)
        self.liner_2=nn.Linear(64,64)
        self.liner_3=nn.Linear(64,1)

    def forward(self,input):
        x=F.relu(self.liner_1(input))
        x = F.relu(self.liner_2(x))
        x=F.sigmoid(self.liner_3(x))
        return x

'''学习效率'''
lr=0.0001

def get_model():
    model=Model()
    opt=torch.optim.Adam(model.parameters(),lr=lr)
    return model,opt
'''计算正确率'''
def accuracy(y_pred,y_label):
    y_pred=(y_pred>0.5).type(torch.int32)
    acc=(y_pred==y_label).float().mean()
    return acc

model,optimizer=get_model()

'''定义损失函数'''
loss_fn=nn.BCELoss()
batch=64
num_of_batch=len(data)//batch
epochs=100

'''使用dataset类进行重构'''
from torch.utils.data import  TensorDataset
from torch.utils.data import  DataLoader


train_x,test_x,train_y,test_y=train_test_split(X,Y)
# '''转化为tensor类型'''
# train_x = torch.from_numpy(train_x).type(torch.FloatTensor)
# test_x = torch.from_numpy(test_x).type(torch.FloatTensor)
# train_y = torch.from_numpy(train_y).type(torch.FloatTensor)
# test_y = torch.from_numpy(test_y).type(torch.FloatTensor)


train_ds=TensorDataset(train_x,train_y)
train_dl=DataLoader(train_ds,batch_size=batch,shuffle=True)
test_ds=TensorDataset(test_x,test_y)
test_dl=DataLoader(test_ds,batch_size=batch)

writer = SummaryWriter('D://tensorboard//example1')
for epoch in range(epochs):
    for x,y in train_dl:
        y_pred=model(x)
        loss=loss_fn(y_pred,y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        epoch_accuracy=accuracy(model(train_x),train_y)
        epoch_loss=loss_fn(model(train_x),train_y).data.item()

        epoch_test_accuracy = accuracy(model(test_x), test_y)
        epoch_test_loss = loss_fn(model(test_x), test_y).data.item()
        '''tensorboard中查看'''
        writer.add_scalar('train_loss',epoch_loss,global_step=epoch)
        writer.add_scalar('test_loss',epoch_test_loss,global_step=epoch)
        writer.add_scalar("test_accuracy",epoch_test_accuracy.item(),global_step=epoch)
        print("epoch: ",epoch,"loss ",round(epoch_loss,3),"mean_accuracy",round(epoch_accuracy.item(),3)
              ,"test_loss",round(epoch_test_loss,3),"test_accuracy",round(epoch_test_accuracy.item(),3))




